#pragma once

#include "buffered_io.h"
#include "class_factory.h"

#include <library/containers/2d_array/2d_array.h>

#include <util/generic/hash_set.h>
#include <util/generic/buffer.h>
#include <util/generic/list.h>
#include <util/generic/maybe.h>
#include <util/generic/bitmap.h>
#include <util/memory/blob.h>
#include <util/digest/murmur.h>

#include <array>
#include <bitset>
#include <list>
#include <string>

#ifdef _MSC_VER
#pragma warning(disable:4127)
#endif


enum ESaverMode
{
    SAVER_MODE_READ        = 1,
    SAVER_MODE_WRITE    = 2,
    SAVER_MODE_WRITE_COMPRESSED = 3,
};

//////////////////////////////////////////////////////////////////////////
struct IBinSaver
{
public:
    typedef unsigned char chunk_id;
    typedef ui32 TStoredSize; // changing this will break compatibility

private:
    // This overload is required to avoid infinite recursion when overriding serialization in derived classes:
    // struct B {
    //  virtual int operator &(IBinSaver& f) {
    //      return 0;
    //  }
    // };
    //
    // struct D : B {
    //  int operator &(IBinSaver& f) override {
    //      f.Add(0, static_cast<B*>(this));
    //      return 0;
    //  }
    // };
    template<class T, typename = decltype(std::declval<T*>()->T::operator& (std::declval<IBinSaver&>()))>
    void CallObjectSerialize(T *p, ui32) { // ui32 will be resolved first if enabled
    // Note: p->operator &(*this) would lead to infinite recursion
        p->T::operator &(*this);
    }

    template<class T, typename = decltype(std::declval<T&>() & std::declval<IBinSaver&>())>
    void CallObjectSerialize(T *p, void*) { // void* will be resolved second if enabled
        (*p) & (*this);
    }

    template<class T>
    void CallObjectSerialize(T *p, ...) { // ... will be resolved last
#if (!defined(_MSC_VER))
        // In MSVC __has_trivial_copy returns false to enums, primitive types and arrays.
        static_assert(__has_trivial_copy(T), "Class is nontrivial copyable, you must define operator&, see");
#endif
        DataChunk(p, sizeof(T));
    }

    // vector
    template <class T, class TA> void DoVector(yvector<T,TA> &data) {
        TStoredSize nSize;
        if (IsReading()) {
            data.clear();
            Add(2, &nSize);
            data.resize(nSize);
        } else {
            nSize = data.size();
            CheckOverflow(nSize, data.size());
            Add(2, &nSize);
        }
        for (TStoredSize i = 0; i < nSize; i++)
            Add(1, &data[i]);
    }

    template <class T, int N> void DoArray(T (&data)[N]) {
        for (size_t i = 0; i < N; i++) {
            Add(1, &(data[i]));
        }
    }

    template <typename TLarge> void CheckOverflow(TStoredSize nSize, TLarge origSize) {
        if (nSize != origSize) {
            fprintf(stderr, "IBinSaver: object size is too large to be serialized (%" PRIu32 " != %" PRIu64 ")\n", nSize, (ui64)origSize);
            abort();
        }
    }

    template <class T, class TA> void DoDataVector(yvector<T,TA> &data) {
        TStoredSize nSize = data.size();
        CheckOverflow(nSize, data.size());
        Add(1, &nSize);
        if (IsReading()) {
            data.clear();
            data.resize(nSize);
        }
        if (nSize > 0)
            DataChunk(&data[0], sizeof(T) * nSize);
    }

    template <class AM> void DoAnyMap(AM& data) {
        if (IsReading()) {
            data.clear();
            TStoredSize nSize;
            Add(3, &nSize);
            yvector<typename AM::key_type, typename AM::allocator_type::template rebind<typename AM::key_type>::other> indices;
            indices.resize(nSize);
            for (TStoredSize i = 0; i < nSize; ++i)
                Add(1, &indices[i]);
            for (TStoredSize i = 0; i < nSize; ++i)
                Add(2, &data[ indices[i] ]);
        } else {
            TStoredSize nSize = data.size();
            CheckOverflow(nSize, data.size());
            Add(3, &nSize);

            yvector<typename AM::key_type, typename AM::allocator_type::template rebind<typename AM::key_type>::other> indices;
            indices.resize(nSize);
            TStoredSize i = 1;
            for (auto pos = data.begin(); pos != data.end(); ++pos, ++i)
                indices[ nSize - i ] = pos->first;
            for (TStoredSize j = 0; j < nSize; ++j)
                Add(1, &indices[j]);
            for (TStoredSize j = 0; j < nSize; ++j)
                Add(2, &data[ indices[j] ]);
        }
    }

    // hash_multimap
    template <class AMM> void DoAnyMultiMap(AMM &data) {
        if (IsReading()) {
            data.clear();
            TStoredSize nSize;
            Add(3, &nSize);
            yvector<typename AMM::key_type, typename AMM::allocator_type::template rebind<typename AMM::key_type>::other> indices;
            indices.resize(nSize);
            for (TStoredSize i = 0; i < nSize; ++i)
                Add(1, &indices[i]);
            for (TStoredSize i = 0; i < nSize; ++i) {
                std::pair<typename AMM::key_type, typename AMM::mapped_type> valToInsert;
                valToInsert.first = indices[i];
                Add(2, &valToInsert.second);
                data.insert(valToInsert);
            }
        } else {
            TStoredSize nSize = data.size();
            CheckOverflow(nSize, data.size());
            Add(3, &nSize);
            for (auto pos = data.begin(); pos != data.end(); ++pos)
                Add(1,(typename AMM::key_type*)(&pos->first));
            for (auto pos = data.begin(); pos != data.end(); ++pos)
                Add(2, &pos->second);
        }
    }

    template <class T> void DoAnySet(T& data) {
        if (IsReading()) {
            data.clear();
            TStoredSize nSize;
            Add(2, &nSize);
            for (TStoredSize i = 0; i < nSize; ++i) {
                typename T::value_type member;
                Add(1, &member);
                data.insert(member);
            }
        } else {
            TStoredSize nSize = data.size();
            CheckOverflow(nSize, data.size());
            Add(2, &nSize);
            for (const auto& elem : data) {
                auto member = elem;
                Add(1, &member);
            }
        }
    }

    // 2D array
    template<class T> void Do2DArray(TArray2D<T> &a) {
        int nXSize = a.GetXSize(), nYSize = a.GetYSize();
        Add(1, &nXSize);
        Add(2, &nYSize);
        if (IsReading())
            a.SetSizes(nXSize, nYSize);
        for (int i = 0; i < nXSize * nYSize; i++)
            Add(3, &a[i/nXSize][i%nXSize]);
    }
    template<class T> void Do2DArrayData(TArray2D<T> &a) {
        int nXSize = a.GetXSize(), nYSize = a.GetYSize();
        Add(1, &nXSize);
        Add(2, &nYSize);
        if (IsReading())
            a.SetSizes(nXSize, nYSize);
        if (nXSize * nYSize > 0)
            DataChunk(&a[0][0], sizeof(T) * nXSize * nYSize);
    }
    // strings
    template<class TStringType> void DataChunkStr(TStringType &data, i64 elemSize)
    {
        if (bRead) {
            TStoredSize nCount = 0;
            File.Read(&nCount, sizeof(TStoredSize));
            data.resize(nCount);
            if (nCount)
                File.Read(data.begin(), nCount * elemSize);
        } else {
            TStoredSize nCount = data.size();
            CheckOverflow(nCount, data.size());
            File.Write(&nCount, sizeof(TStoredSize));
            File.Write(data.c_str(), nCount * elemSize);
        }
    }
    void DataChunkString(std::string &data)
    {
        DataChunkStr(data, sizeof(char));
    }
    void DataChunkStroka(TString &data)
    {
        DataChunkStr(data, sizeof(TString::char_type));
    }
    void DataChunkWtroka(TUtf16String &data)
    {
        DataChunkStr(data, sizeof(TUtf16String::char_type));
    }

    void DataChunk(void *pData, i64 nSize)
    {
        i64 chunkSize = 1 << 30;
        for (i64 offset = 0; offset < nSize; offset += chunkSize) {
            void *ptr = (char *)pData + offset;
            i64 size = offset + chunkSize < nSize ? chunkSize : (nSize - offset);
            if (bRead)
                File.Read(ptr, size);
            else
                File.Write(ptr, size);
        }
    }

    // storing/loading pointers to objects
    void StoreObject(TObjectBase *pObject);
    TObjectBase* LoadObject();

    bool bRead;
    TBufferedStream<> File;
    // maps objects addresses during save(first) to addresses during load(second) - during loading
    // or serves as a sign that some object has been already stored - during storing
    bool StableOutput;

    typedef yhash<void*, ui32, TPtrHash> PtrIdHash;
    TAutoPtr<PtrIdHash> PtrIds;

    typedef yhash<ui64, TPtr<TObjectBase> > CObjectsHash;
    TAutoPtr<CObjectsHash> Objects;

    yvector<TObjectBase*> ObjectQueue;
public:
    bool IsReading() { return bRead; }
    void AddRawData(const chunk_id, void *pData, i64 nSize) { DataChunk(pData, nSize); }

    // return type of Add() is used to detect specialized serializer (see HasNonTrivialSerializer below)
    template<class T>
        char Add(const chunk_id, T *p)
    {
        CallObjectSerialize(p, 0u);
        return 0;
    }
    int Add(const chunk_id, std::string *pStr)
    {
        DataChunkString(*pStr);
        return 0;
    }
    int Add(const chunk_id, TString *pStr)
    {
        DataChunkStroka(*pStr);
        return 0;
    }
    int Add(const chunk_id, TUtf16String *pStr)
    {
        DataChunkWtroka(*pStr);
        return 0;
    }
    int Add(const chunk_id, TBlob *blob)
    {
        if (bRead) {
            ui64 size = 0;
            File.Read(&size, sizeof(size));
            TBuffer buffer;
            buffer.Advance(size);
            if (size > 0)
                File.Read(buffer.Data(), buffer.Size());
            (*blob) = TBlob::FromBuffer(buffer);
        } else {
            const ui64 size = blob->Size();
            File.Write(&size, sizeof(size));
            File.Write(blob->Data(), blob->Size());
        }
        return 0;
    }
    template<class T1, class TA>
        int Add(const chunk_id, yvector<T1,TA> *pVec)
    {
        if (HasNonTrivialSerializer<T1>(0u))
            DoVector(*pVec);
        else
            DoDataVector(*pVec);
        return 0;
    }

    template<class T, int N>
        int Add(const chunk_id, T (*pVec)[N])
    {
        if (HasNonTrivialSerializer<T>(0u))
            DoArray(*pVec);
        else
            DataChunk(pVec, sizeof(*pVec));
        return 0;
    }

    template<class T1, class T2, class T3, class T4>
        int Add(const chunk_id, ymap<T1,T2,T3,T4> *pMap)
    {
        DoAnyMap(*pMap);
        return 0;
    }
    template<class T1, class T2, class T3, class T4, class T5>
        int Add(const chunk_id, yhash<T1,T2,T3,T4,T5> *pHash)
    {
        DoAnyMap(*pHash);
        return 0;
    }
    template<class T1, class T2, class T3, class T4, class T5>
        int Add(const chunk_id, yhash_mm<T1,T2,T3,T4,T5> *pHash)
    {
        DoAnyMultiMap(*pHash);
        return 0;
    }
    template<class K, class L, class A>
        int Add(const chunk_id, yset<K, L, A> *pSet)
    {
        DoAnySet(*pSet);
        return 0;
    }
    template<class T1, class T2, class T3, class T4>
        int Add(const chunk_id, yhash_set<T1,T2,T3,T4> *pHash)
    {
        DoAnySet(*pHash);
        return 0;
    }

    template<class T1>
        int Add(const chunk_id, TArray2D<T1> *pArr)
    {
        if (HasNonTrivialSerializer<T1>(0u))
            Do2DArray(*pArr);
        else
            Do2DArrayData(*pArr);
        return 0;
    }
    template<class T1>
        int Add(const chunk_id, ylist<T1> *pList)
    {
        ylist<T1> &data = *pList;
        if (IsReading())
        {
            int nSize;
            Add(2, &nSize);
            data.clear();
            data.insert(data.begin(), nSize, T1());
        }
        else
        {
            int nSize = data.size();
            Add(2, &nSize);
        }
        int i = 1;
        for (typename ylist<T1>::iterator k = data.begin(); k != data.end(); ++k, ++i)
            Add(i + 2, &(*k));
        return 0;
    }
    template <class T1, class T2>
        int Add(const chunk_id, std::pair<T1, T2> *pData)
    {
        Add(1, &(pData->first));
        Add(2, &(pData->second));
        return 0;
    }

    template <class T1, size_t N>
    int Add(const chunk_id, std::array<T1, N> *pData)
    {
        if (HasNonTrivialSerializer<T1>(0u)) {
            for (size_t i = 0; i < N; ++i)
                Add(1, &(*pData)[i]);
        } else {
            DataChunk((void*)pData->data(), pData->size()*sizeof(T1));
        }
        return 0;
    }

    template <size_t N>
    int Add(const chunk_id, std::bitset<N> *pData)
    {
        if (IsReading()) {
            std::string s;
            Add(1, &s);
            *pData = std::bitset<N>(s);
        } else {
            std::string s = pData->template to_string< char, std::char_traits<char>, std::allocator<char> >();
            Add(1, &s);
        }
        return 0;
    }

    int Add(const chunk_id, TDynBitMap* pData)
    {
        if (IsReading()) {
            ui64 count = 0;
            Add(1, &count);
            pData->Clear();
            pData->Reserve(count * sizeof(TDynBitMap::TChunk) * 8);
            for (ui64 i = 0; i < count; ++i) {
                TDynBitMap::TChunk chunk = 0;
                Add(i + 1, &chunk);
                if (i > 0) {
                    pData->LShift(8 * sizeof(TDynBitMap::TChunk));
                }
                pData->Or(chunk);
            }
        } else {
            ui64 count = pData->GetChunkCount();
            Add(1, &count);
            for (ui64 i = 0; i < count; ++i) {
                // Write in reverse order
                TDynBitMap::TChunk chunk = pData->GetChunks()[count - i - 1];
                Add(i + 1, &chunk);
            }
        }
        return 0;
    }

    void AddPolymorphicBase(chunk_id, TObjectBase *pObject) {
        (*pObject) & (*this);
    }

    template <class T1, class T2>
    void DoPtr(TPtrBase<T1,T2> *pData) {
        if (pData && pData->Get()) {
        }
        if (IsReading())
            pData->Set(CastToUserObject(LoadObject(), (T1*)nullptr));
        else
            StoreObject(pData->GetBarePtr());
    }
    template <class T>
        int Add(const chunk_id, TMaybe<T> *pData)
    {
        TMaybe<T>& data = *pData;
        if (IsReading()) {
            bool defined = false;
            Add(1, &defined);
            if (defined) {
                data = T();
                Add(2, data.Get());
            }
        } else {
            bool defined = data.Defined();
            Add(1, &defined);
            if (defined) {
                Add(2, data.Get());
            }
        }
        return 0;
    }

    template <typename TOne>
    void AddMulti(TOne &one)
    {
        Add(0, &one);
    }

    template <typename THead, typename ... TTail>
        void AddMulti(THead &head, TTail& ... tail)
    {
        Add(0, &head);
        AddMulti(tail ...);
    }

    template<class T, typename = decltype(std::declval<T&>() & std::declval<IBinSaver&>())>
    static bool HasNonTrivialSerializer(ui32) {
        return true;
    }

    template<class T>
    static bool HasNonTrivialSerializer(...) {
        return sizeof(std::declval<IBinSaver*>()->Add(0, std::declval<T*>())) != 1;
    }
public:
    IBinSaver(IBinaryStream &stream, bool _bRead, bool stableOutput = false) : bRead(_bRead), File(_bRead, stream), StableOutput(stableOutput) {}
    virtual ~IBinSaver();
    bool IsValid() const { return File.IsValid(); }
};

// realisation of forward declared serialisation operator
template< class TUserObj, class TRef>
int TPtrBase<TUserObj,TRef>::operator&(IBinSaver &f) {
    f.DoPtr(this);
    return 0;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

extern TClassFactory<TObjectBase> *pSaverClasses;
void StartRegisterSaveload();

template<class TReg>
struct TRegisterSaveLoadType {
    TRegisterSaveLoadType(int num) {
        StartRegisterSaveload();
        pSaverClasses->RegisterType(num, TReg::NewSaveLoadNullItem, (TReg*)nullptr);
    }
};

#define Y_BINSAVER_REGISTER(name)\
    BASIC_REGISTER_CLASS(name)\
    static TRegisterSaveLoadType<name> init##name(MurmurHash<int>(#name, sizeof(#name)));

#define REGISTER_SAVELOAD_CLASS(N, name)\
    BASIC_REGISTER_CLASS(name)\
    static TRegisterSaveLoadType<name> init##name##N(N);

// using TObj/TRef on forward declared templ class will not work
// but multiple registration with same id is allowed
#define REGISTER_SAVELOAD_TEMPL1_CLASS(N, className, T)\
    static TRegisterSaveLoadType< className<T> > init##className##T##N(N);

#define REGISTER_SAVELOAD_TEMPL2_CLASS(N, className, T1, T2)\
    typedef className<T1,T2> temp##className##T1##_##T2##temp;\
    static TRegisterSaveLoadType< className<T1,T2> > init##className##T1##_##T2##N(N);

#define REGISTER_SAVELOAD_TEMPL3_CLASS(N, className, T1, T2, T3)\
    typedef className<T1,T2,T3> temp##className##T1##_##T2##_##T3##temp;\
    static TRegisterSaveLoadType< className<T1,T2,T3> > init##className##T1##_##T2##_##T3##N(N);

#define REGISTER_SAVELOAD_NM_CLASS(N, nmspace, className)\
    BASIC_REGISTER_CLASS(nmspace::className)\
    static TRegisterSaveLoadType<nmspace::className> init_##nmspace##_##name##N(N);

#define REGISTER_SAVELOAD_TEMPL1_NM_CLASS(N, nmspace, className, T)\
    typedef nmspace::className<T> temp_init##nmspace##className##T##temp;\
    BASIC_REGISTER_CLASS(nmspace::className<T>)\
    static TRegisterSaveLoadType< nmspace::className<T> > temp_init##nmspace##_##name##T##N(N);


#define REGISTER_SAVELOAD_CLASS_NAME(N, cls, name) \
    BASIC_REGISTER_CLASS(cls) \
    static TRegisterSaveLoadType< cls > init ## name ## N(N);

#define REGISTER_SAVELOAD_CLASS_NS_PREF(N, cls, ns, pref) \
    REGISTER_SAVELOAD_CLASS_NAME(N, ns :: cls, _ ## pref ## _ ## cls)

#define SAVELOAD(...) int operator& (IBinSaver &f) { f.AddMulti(__VA_ARGS__); return 0; }

#define SAVELOAD_OVERRIDE(base, ...) \
    int operator& (IBinSaver &f) override { \
        base::operator& (f); \
        f.AddMulti(__VA_ARGS__); \
        return 0; \
    }

#define SAVELOAD_BASE(...) \
    int operator& (IBinSaver &f) { \
        TBase::operator& (f); \
        f.AddMulti(__VA_ARGS__); \
        return 0; \
    }
